package com.bzframework.pro.socket;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.bzframework.pro.authorize.constants.TokenNameEnum;
import com.bzframework.pro.authorize.token.UserToken;
import com.bzframework.pro.authorize.token.manager.UserTokenManager;
import com.bzframework.pro.socket.manager.MessagingManager;
import com.bzframework.pro.socket.message.Message;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.util.StringUtils;
import org.springframework.web.reactive.socket.CloseStatus;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.Disposable;
import reactor.core.publisher.BufferOverflowStrategy;
import reactor.core.publisher.Mono;
import reactor.core.publisher.SignalType;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@AllArgsConstructor
@Slf4j
public class WebSocketMessagingHandler implements WebSocketHandler {

    private final MessagingManager messagingManager;

    private final UserTokenManager userTokenManager;

    private static final Integer MAX_BUFFER_SIZE = Integer.getInteger("bzframework.websocket.subscriber.buffer-size", 1024);

    @Override
    @SuppressWarnings("all")
    public Mono<Void> handle(WebSocketSession session) {

        // websocket 当前会话所有订阅信息
        Map<String, Disposable> subs = new ConcurrentHashMap<>();

        return hasAuthentication(session)
                .switchIfEmpty(session.send(Mono.just(session.textMessage(JSON.toJSONString(Message.authError()))))
                                      .then(session.close(CloseStatus.BAD_DATA))
                                      .then(Mono.empty()))
                .flatMap(userToken -> {
                    // ws 连接建立后 发送消息回调
                    return session
                            .receive()
                            .doOnNext(message -> {

                                if (message.getType() == WebSocketMessage.Type.PONG) {
                                    return;
                                }
                                if (message.getType() == WebSocketMessage.Type.PING) {
                                    session.send(Mono.just(session.pongMessage(DataBufferFactory::allocateBuffer)))
                                           .subscribe();
                                    return;
                                }
                                // 解析request请求
                                MessagingRequest request = JSON.parseObject(message.getPayloadAsText(), MessagingRequest.class);
                                if (request == null) {
                                    return;
                                }
                                if (!StringUtils.hasText(request.getId())) {
                                    sendToSession(session, Message.error(request.getType().name(), null, "id不能为空"));
                                    return;
                                }
                                // payload ping -> pong
                                if (request.getType() == MessagingRequest.Type.ping) {
                                    sendToSession(session, Message.pong(request.getId()));
                                    return;
                                }
                                if (request.getType() == MessagingRequest.Type.sub) {

                                    Disposable old = subs.get(request.getId());
                                    if (old != null && !old.isDisposed()) {
                                        return;
                                    }

                                    Disposable subscription = messagingManager
                                            .subscribe(request)
                                            .onBackpressureBuffer(MAX_BUFFER_SIZE, BufferOverflowStrategy.DROP_OLDEST)
                                            .onErrorResume(err -> Mono.just(Message.error(request.getId(), request.getTopic(), err)))
                                            .map(msg -> session.textMessage(JSON.toJSONString(msg)))
                                            .flatMap(msg -> session.send(Mono.just(msg)))
                                            .contextWrite(ctx -> {
                                                return ctx.put(UserToken.class, userToken);
                                            })
                                            .doFinally(signalType -> {

                                                if (signalType == SignalType.ON_COMPLETE) {
                                                    Mono.just(session.textMessage(JSON.toJSONString(
                                                                Message.complete(request.getId()),
                                                                SerializerFeature.DisableCircularReferenceDetect)))
                                                        .as(session::send)
                                                        .subscribe();
                                                } else {
                                                    log.debug("{} subscription:{}", signalType, request.getTopic());
                                                }

                                            })
                                            .subscribe();
                                    if (!subscription.isDisposed()) {
                                        subs.put(request.getId(), subscription);
                                    }
                                }else if (request.getType() == MessagingRequest.Type.unsub) {
                                    subs.remove(request.getId()).dispose();
                                }else {
                                    sendToSession(session, Message.error(request.getId(), request.getTopic(), "不支持的类型:" + request.getType()));
                                }
                            })
                            .then();
                }).doFinally(signalType -> {
                    // 连接断开时，释放全部订阅
                    subs.values().forEach(Disposable::dispose);
                    subs.clear();
                });
    }


    private void sendToSession(WebSocketSession session, Message message) {
        session.send(Mono.just(session.textMessage(JSON.toJSONString(message))))
               .subscribe();
    }


    private Mono<UserToken> hasAuthentication(WebSocketSession session) {

        String token = null;
        // 1.header
        if (session.getHandshakeInfo().getHeaders().containsKey(TokenNameEnum.HEADER_TOKEN_NAME.getText())) {
            token = session.getHandshakeInfo().getHeaders().getFirst(TokenNameEnum.HEADER_TOKEN_NAME.getText());
        } else {
            // 2.path
            String paths = session.getHandshakeInfo().getUri().getPath();
            String[] path = paths.split("[/]");
            if (path.length == 0) {
                return Mono.empty();
            }
            token = path[path.length - 1];
        }
        if (!StringUtils.hasText(token)) {
            return Mono.empty();
        }
        return userTokenManager.getUserToken(token);
    }


}
